import cv2
import numpy as np
from PyQt5.QtWidgets import QWidget, QLabel, QGridLayout, QVBoxLayout, QHBoxLayout
from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen, QColor
from PyQt5.QtCore import Qt, QRect, pyqtSignal

class GridViewer(QWidget):
    """
    网格查看器组件，用于显示和交互式编辑网格
    """
    grid_selected = pyqtSignal(int, int)  # 发射网格选择信号 (行, 列)
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.image = None
        self.grid_size = 30
        self.selected_row = -1
        self.selected_col = -1
        self.setMouseTracking(True)
        self.init_ui()
    
    def init_ui(self):
        self.layout = QVBoxLayout()
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignCenter)
        self.layout.addWidget(self.image_label)
        self.setLayout(self.layout)
    
    def set_image(self, image, grid_size):
        """
        设置图像和网格大小
        
        Args:
            image: 输入图像 (numpy数组，RGB格式)
            grid_size: 网格大小
        """
        self.image = image.copy()
        self.grid_size = grid_size
        self.update_display()
    
    def update_display(self):
        """
        更新显示
        """
        if self.image is None:
            return
        
        # 创建带网格的图像
        display_image = self.image.copy()
        h, w = display_image.shape[:2]
        
        # 绘制网格线
        for x in range(0, w, self.grid_size):
            cv2.line(display_image, (x, 0), (x, h), (0, 0, 0), 1)
        
        for y in range(0, h, self.grid_size):
            cv2.line(display_image, (0, y), (w, y), (0, 0, 0), 1)
        
        # 如果有选中的网格，高亮显示
        if self.selected_row >= 0 and self.selected_col >= 0:
            y1 = self.selected_row * self.grid_size
            x1 = self.selected_col * self.grid_size
            y2 = min(y1 + self.grid_size, h)
            x2 = min(x1 + self.grid_size, w)
            
            # 绘制高亮边框
            cv2.rectangle(display_image, (x1, y1), (x2, y2), (255, 0, 0), 2)
        
        # 转换为QImage并显示
        h, w, c = display_image.shape
        bytes_per_line = 3 * w
        q_image = QImage(display_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(q_image)
        
        # 调整大小以适应标签
        label_size = self.image_label.size()
        pixmap = pixmap.scaled(label_size, Qt.KeepAspectRatio, Qt.SmoothTransformation)
        
        self.image_label.setPixmap(pixmap)
    
    def mousePressEvent(self, event):
        """
        鼠标点击事件处理
        """
        if self.image is None:
            return
        
        # 获取图像在标签中的位置和大小
        pixmap = self.image_label.pixmap()
        if pixmap is None:
            return
            
        img_rect = self.image_label.contentsRect()
        scaled_width = pixmap.width()
        scaled_height = pixmap.height()
        
        # 计算图像在标签中的实际位置
        x_offset = (img_rect.width() - scaled_width) // 2
        y_offset = (img_rect.height() - scaled_height) // 2
        
        # 计算点击位置在图像中的坐标
        x = event.x() - self.image_label.x() - x_offset
        y = event.y() - self.image_label.y() - y_offset
        
        # 计算缩放比例
        scale_x = self.image.shape[1] / scaled_width
        scale_y = self.image.shape[0] / scaled_height
        
        # 转换为原始图像坐标
        orig_x = int(x * scale_x)
        orig_y = int(y * scale_y)
        
        # 确保坐标在图像范围内
        if 0 <= orig_x < self.image.shape[1] and 0 <= orig_y < self.image.shape[0]:
            # 计算网格行列
            row = orig_y // self.grid_size
            col = orig_x // self.grid_size
            
            # 更新选中的网格
            self.selected_row = row
            self.selected_col = col
            
            # 发射信号
            self.grid_selected.emit(row, col)
            
            # 更新显示
            self.update_display()

class FeatureVisualizer(QWidget):
    """
    特征可视化组件
    """
    def __init__(self, parent=None):
        super().__init__(parent)
        self.features = None
        self.grid_size = 30
        self.init_ui()
    
    def init_ui(self):
        self.layout = QVBoxLayout()
        
        # 特征显示标签
        self.feature_label = QLabel("特征可视化")
        self.feature_label.setAlignment(Qt.AlignCenter)
        self.feature_label.setMinimumSize(400, 300)
        self.feature_label.setStyleSheet("border: 1px solid black")
        self.layout.addWidget(self.feature_label)
        
        # 特征详情区域
        self.details_layout = QGridLayout()
        
        # 颜色特征
        self.color_label = QLabel("颜色特征:")
        self.color_value = QLabel("")
        self.details_layout.addWidget(self.color_label, 0, 0)
        self.details_layout.addWidget(self.color_value, 0, 1)
        
        # 纹理特征
        self.texture_label = QLabel("纹理特征:")
        self.texture_value = QLabel("")
        self.details_layout.addWidget(self.texture_label, 1, 0)
        self.details_layout.addWidget(self.texture_value, 1, 1)
        
        # 形状特征
        self.shape_label = QLabel("形状特征:")
        self.shape_value = QLabel("")
        self.details_layout.addWidget(self.shape_label, 2, 0)
        self.details_layout.addWidget(self.shape_value, 2, 1)
        
        self.layout.addLayout(self.details_layout)
        self.setLayout(self.layout)
    
    def set_features(self, features, grid_size):
        """
        设置特征数据
        
        Args:
            features: 特征字典，键为(行,列)，值为特征向量
            grid_size: 网格大小
        """
        self.features = features
        self.grid_size = grid_size
    
    def show_grid_features(self, row, col):
        """
        显示指定网格的特征
        
        Args:
            row: 网格行索引
            col: 网格列索引
        """
        if self.features is None or (row, col) not in self.features:
            self.color_value.setText("无数据")
            self.texture_value.setText("无数据")
            self.shape_value.setText("无数据")
            return
        
        # 获取特征向量
        feature = self.features[(row, col)]
        
        # 显示颜色特征 (前6个值)
        color_text = f"H均值: {feature[0]:.2f}, H方差: {feature[1]:.2f}\n"
        color_text += f"S均值: {feature[2]:.2f}, S方差: {feature[3]:.2f}\n"
        color_text += f"V均值: {feature[4]:.2f}, V方差: {feature[5]:.2f}"
        self.color_value.setText(color_text)
        
        # 显示纹理特征 (中间部分)
        texture_idx = 6 + 8 + 4 + 4  # 颜色特征后的索引
        texture_text = f"对比度: {np.mean(feature[texture_idx:texture_idx+4]):.4f}\n"
        texture_text += f"相异性: {np.mean(feature[texture_idx+4:texture_idx+8]):.4f}\n"
        texture_text += f"同质性: {np.mean(feature[texture_idx+8:texture_idx+12]):.4f}"
        self.texture_value.setText(texture_text)
        
        # 显示形状特征 (最后部分)
        shape_idx = texture_idx + 20  # 纹理特征后的索引
        shape_text = f"边缘密度: {feature[shape_idx]:.4f}\n"
        shape_text += f"主方向: {np.argmax(feature[shape_idx+1:]) * 30}°"
        self.shape_value.setText(shape_text)
        
    def create_feature_visualization(self, image, features, grid_size):
        """
        创建特征可视化图像
        
        Args:
            image: 原始图像
            features: 特征字典
            grid_size: 网格大小
            
        Returns:
            可视化图像
        """
        h, w = image.shape[:2]
        visualization = np.zeros((h, w, 3), dtype=np.uint8)
        
        # 为每个网格创建颜色编码
        for (row, col), feature in features.items():
            y = row * grid_size
            x = col * grid_size
            end_y = min(y + grid_size, h)
            end_x = min(x + grid_size, w)
            
            # 使用HSV颜色特征来可视化
            h_val = int(feature[0] / 180 * 255)  # H值
            s_val = int(feature[2])              # S值
            v_val = int(feature[4])              # V值
            
            # 填充网格
            visualization[y:end_y, x:end_x] = [h_val, s_val, v_val]
        
        # 添加网格线
        for x in range(0, w, grid_size):
            cv2.line(visualization, (x, 0), (x, h), (255, 255, 255), 1)
        
        for y in range(0, h, grid_size):
            cv2.line(visualization, (0, y), (w, y), (255, 255, 255), 1)
        
        return visualization
    
    def update_visualization(self, image, features, grid_size):
        """
        更新特征可视化
        
        Args:
            image: 原始图像
            features: 特征字典
            grid_size: 网格大小
        """
        self.set_features(features, grid_size)
        
        # 创建可视化图像
        viz_image = self.create_feature_visualization(image, features, grid_size)
        
        # 转换为QImage并显示
        h, w, c = viz_image.shape
        bytes_per_line = 3 * w
        q_image = QImage(viz_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(q_image)
        
        # 调整大小以适应标签
        label_size = self.feature_label.size()
        pixmap = pixmap.scaled(label_size, Qt.KeepAspectRatio, Qt.SmoothTransformation)
        
        self.feature_label.setPixmap(pixmap)